import decimal
import json
import logging

from app.migrate_neo4j.base import BaseMigrator

logger = logging.getLogger(__name__)


class MockLLMDataModeler:
    def get_ddl_for_tables(self, inspector):
        ddl_strings = []
        for table_name in inspector.get_table_names():
            ddl = f"CREATE TABLE {table_name} (\n"
            for col in inspector.get_columns(table_name):
                ddl += f"  {col['name']} {col['type']},\n"
            ddl += ");"
            ddl_strings.append(ddl)
        return "\n\n".join(ddl_strings)

    def generate_model(self, schema_ddl):
        logging.info("[LLM] Generating graph model from DDL...")
        mock_response = """
        {
          "nodes": {
            "customers": {"label": "Customer"}, "products": {"label": "Product"},
            "orders": {"label": "Order"}, "order_items": {"is_relationship_table": true}
          },
          "relationships": [
            {"from_table": "orders", "to_table": "customers", "type": "PLACED_BY"},
            {"from_table": "order_items", "to_table": "orders", "type": "PART_OF"},
            {"from_table": "order_items", "to_table": "products", "type": "INCLUDES"}
          ]
        }
        """
        return json.loads(mock_response)


# -----------------------------------------------------------------------------
# ---  基于 LLM 建模的 Migrator ---
# -----------------------------------------------------------------------------
class LLMMigrator(BaseMigrator):
    """
    将数据库 DDL 语句发送给 LLM，让其进行建模。
    适用于数据库设计不严谨或需要更智能语义关系的场景。
    """

    def __init__(self, rdbms_config, neo4j_config, batch_size=1000):
        super().__init__(rdbms_config, neo4j_config, batch_size)
        self.llm_modeler = MockLLMDataModeler()

    def model_schema(self):
        logging.info("[MODEL] Modeling schema using LLM.")
        ddl = self.llm_modeler.get_ddl_for_tables(self.inspector)
        llm_model = self.llm_modeler.generate_model(ddl)

        # 将 LLM 的简洁模型转换为脚本所需的完整模型
        model = {"nodes": {}, "relationships": []}
        for table, info in llm_model["nodes"].items():
            pk = self.inspector.get_pk_constraint(table)['constrained_columns'][0]
            model["nodes"][table] = {"label_name": info["label"], "primary_key": pk, **info}

        for rel in llm_model["relationships"]:
            from_table, to_table = rel["from_table"], rel["to_table"]
            # 找到定义这个关系的外键
            fks = self.inspector.get_foreign_keys(from_table)
            fk_info = next((fk for fk in fks if fk['referred_table'] == to_table), None)
            if not fk_info:  # 检查反向
                fks = self.inspector.get_foreign_keys(to_table)
                fk_info = next((fk for fk in fks if fk['referred_table'] == from_table), None)
                if fk_info: from_table, to_table = to_table, from_table  # 反转

            if fk_info:
                model["relationships"].append({
                    "source_table_name": from_table,
                    "start_node_label": model["nodes"][to_table]["label_name"],
                    "start_node_key": fk_info['referred_columns'][0],
                    "end_node_label": model["nodes"][from_table]["label_name"],
                    "end_node_key": model["nodes"][from_table]["primary_key"],
                    "rel_type": rel["type"],
                    "fk_column": fk_info['constrained_columns'][0]
                })
        self.graph_model = model

    def transform_data(self, data_row, context):
        if context['type'] == 'node':
            for key, value in data_row.items():
                if isinstance(value, decimal.Decimal):
                    data_row[key] = float(value)
            return data_row
        elif context['type'] == 'relationship':
            fk_val = data_row.get(context['fk_column'])
            pk_val = data_row.get(context['end_node_key'])
            if fk_val is not None and pk_val is not None:
                return {"start_id": fk_val, "end_id": pk_val}
        return None
